/*
    Copyright 2006-2011 Patrik Jonsson, sunrise@familjenjonsson.org

    This file is part of Sunrise.

    Sunrise is free software; you can redistribute it and/or modify
    it under the terms of the GNU General Public License as published by
    the Free Software Foundation; either version 3 of the License, or
    (at your option) any later version.

    Sunrise is distributed in the hope that it will be useful,
    but WITHOUT ANY WARRANTY; without even the implied warranty of
    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
    GNU General Public License for more details.

    You should have received a copy of the GNU General Public License
    along with Sunrise.  If not, see <http://www.gnu.org/licenses/>.

*/

/// \file
/// Contains the function for doing the iterative calculation of the
/// equilibrium radiation field in the presence of self-absorption.

// $Id$

#ifndef __equilibrium__
#define __equilibrium__

#include <iostream> 
#include <vector> 
#include "xfer.h"
#include "shoot.h"
#include "full_sed_emergence.h"
#include "boost/lexical_cast.hpp"
#include "mcrx-debug.h"

namespace mcrx {
  template<typename T_dust_model, typename T_grid, typename T_emission >
  bool
  determine_dust_equilibrium_intensities
  (
    T_dust_model& model,
    T_grid& grid,
    T_emission& emission,
    const typename T_dust_model::T_lambda& lambda,
    const array_2& heating_intensity,
    array_2& dust_intensity,
    /** The state vector which is used to initialize the
	random number generators in the threads.  Must
	contain n_threads entries. */
    std::vector<T_rng::T_state>& rng_states,
    long n_rays_desired,
    int n_threads,
    bool bind_threads,
    int bank_size,
    T_float tolerance,
    T_float i_ls,
    int n_forced,
    T_float i_s,
    scatter_shooter<typename T_dust_model::T_biaser> shooter,
    const terminator& term,
    bool check_energy
    );
 }


/** Iterates to convergence to determine the dust temperature
    distribution when self-absorption is included. This involves
    repeatedly calculating the dust temperature, how much extra dust
    emission this implies, and then running a ray tracing step to
    determine how this emission in turn heats other dust. The output
    is the dust_intensity array, which contains the radiation
    intensity due to dust emission in the grid cells.  Memory use can
    be very large in this function, because there are several arrays
    that are of size number of cells * number of wavelengths. For this
    reason, care must be taken to minimize temporaries. */
template<typename T_dust_model, typename T_grid, typename T_emission >
bool
mcrx::determine_dust_equilibrium_intensities
(
 /** The dust model which keeps the scatterers. Note that the dust
     model must be a class derived from grain_model so that it has
     functionality for calculating grain temperatures. */
 T_dust_model& model,
 T_grid& grid,
 T_emission& emission,
 /** The wavelength vector used. Note that it is assumed that the dust
     model has already been set to this wavelength vector. */
 const typename T_dust_model::T_lambda& lambda,
 /** The radiation intensity in the grid cells determined by the
     transfer of the primary (stellar) emission. */
 const array_2& heating_intensity,
 /** The output radiation intensity in the cells due to dust emission,
     after the equilibrium has been determined. */
 array_2& dust_intensity,
 /** The state vector which is used to initialize the
     random number generators in the threads.  Must
     contain n_threads entries. */
 std::vector<T_rng::T_state>& rng_states,
 /// The number of rays to shoot in the ray tracing stages.
 long n_rays_desired,
 /// The number of threads to use.
 int n_threads,
 /// If true, the threads are bound to the cpus.
 bool bind_threads,
 /// The thread bank size,
 int bank_size,
 /** The tolerance for determining when the iteration has
     converged. Since the calculation works by emitting the delta of
     the cell luminosities in each iteration, the luminosity in the
     grid will decrease over time as luminosity escapes the
     volume. When the dust luminosity in the grid has decreased to
     tolerance*L_initial, the iteration stops. */
 T_float tolerance,
 /** The intensity below which the ray is subject to Russian Roulette. */
 T_float i_ls,
 /** The number of forced scatterings used when shooting. While
     scattering is generally not important in the IR, using it also
     means a more accuarate calculation of the radiative intensity. */ 
 int n_forced,
 /** The intensity above which the rays are split. */
 T_float i_s,
 /// Shooter propagates necessary parameters to the shoot() function.
 scatter_shooter<typename T_dust_model::T_biaser> shooter,
 const terminator& term,
 bool check_energy
 )
{
  using namespace std;
  using namespace blitz;

  typedef typename T_dust_model::T_lambda T_lambda;
  typedef typename T_dust_model::T_biaser T_biaser;
  typedef scatter_shooter<T_biaser> T_shooter;
  typedef emergence<T_lambda> T_emergence;

  ASSERT_ALL(heating_intensity>=0);
  ASSERT_ALL(heating_intensity<1e300);

  const T_float area_factor = grid.area_factor();

  cout << "Calculating dust equilibrium temperatures\n";

  // This vector keeps the same pointers to scatterers as the dust
  // model, but these are grain_model* and are used for the SED
  // calculation. 
  std::vector<grain_model<polychromatic_scatterer_policy, 
    mcrx_rng_policy>*> grain_models;
  for(int i=0; i< model.n_scatterers(); ++i) {
    grain_models.push_back
      (dynamic_cast<grain_model<polychromatic_scatterer_policy, 
       mcrx_rng_policy>*>(&model.get_scatterer(i)));
  }

  // We use an empty emergence object, since when just calculating
  // the intensities we can save time by not using any cameras.
  T_emergence eme (vector<std::pair<T_float, T_float> >(),
		   1,1,1);

  // check that output dust_intensity array is right size and initialize
  ASSERT_ALL(dust_intensity.shape()==heating_intensity.shape());
  dust_intensity = 0;

  const int nc = grid.n_cells();

  // save the volumes of the cells and the opacities for use in
  // convergence calculations
  array_1 volume(nc);

  // save the densities of the grid cells in a compact array
  array_2 densities(nc, model.zero_density().size());
  {
    int c=0;
    for (typename T_grid::const_iterator cc=grid.begin(); 
	 cc!=grid.end(); ++cc, ++c) {
      volume(c)= cc->volume();
      densities(c, Range::all())= 
	cc->data()->get_absorber().densities();
    }
  }

  // Calculate the amount of absorbed luminosity due to primary heating.
  array_1 orig_abs(nc);
  for(int c=0; c<nc; ++c) {
    orig_abs(c) = 
      integrate_quantity(heating_intensity(c, Range::all())*
			 model.total_abs_length_to_absorption(densities(c, Range::all()))*
			 4*constants::pi*volume(c)*area_factor,
			 lambda, false);
  }

  bool converged = false;
  T_float grid_absorption = 0;
  T_float L_initial;
  int iter=1;
  while(!converged) {
    cout << "\nIteration " << iter << endl;
     
    // The dust SED calculation relies on the intensities being in
    // physical units, not simulation units, so we must take care to
    // convert dust_intensity to the same units as heating_intensity.
    // For all iterations except first, use perturbative calc.
     
    // Ask the emission object (the ir_grid) to calculate the emission
    // SEDs of the cells based on the current estimate of the
    // radiation field. We pass the radiation field, which is the sum
    // of primary heating and the current estimate of the dust
    // intensity as an expression template to avoid allocating a large
    // temporary array.
    emission.calculate_SED(heating_intensity+dust_intensity, 
			   grain_models, term, n_threads, 0,
			   (iter>1));
     
    const T_float current_lum = emission.total_emission_weight();
    cout << "Luminosity in this iteration: " <<  current_lum << endl;
    if (iter==1) {
      L_initial=current_lum;

      if(L_initial==0) {
	cout << "No luminosity in grid, nothing to calculate." << endl;
	converged = true;
	continue;
      }
    }

    // If we have (apparent) energy nonconservation during the SED
    // calculation, print a warning. This should NEVER happen.
    if((iter>1)&&(current_lum > grid_absorption*1.01)) {
      cerr << "Fatal: luminosity nonconservation in SED calculation by " << current_lum/grid_absorption-1 <<  "\n" << endl;
      throw 0;
    }

    // now shoot so we get the updated intensities. We must reset the
    // grid first so the old intensities are erased.
    cout << "Shooting" << endl;
    grid.reset();
    model.reset();
    long n_rays=0;
     
    // START SHOOT
    assert (n_threads == rng_states.size()); 
     
    // The xfer rngs are loaded in the thread shooting function.
    mcrx_rng_policy rng;
     
    xfer<T_dust_model, T_grid, mcrx_rng_policy> x (grid, emission, eme, model,
						   rng, i_ls, 0, n_forced, i_s, 
						   false);
     
    if( shoot (x, shooter, term, rng_states, 
	       n_rays_desired, n_rays, n_threads, bind_threads,
	       bank_size, 11, "shoot in determine dust equilibrium") )
      // we got termination, return true and stop
      // we are in a consistent state, so we can return here
      return true;
    // END SHOOT
     
    const T_float normalization = emission.total_emission_weight()/n_rays;

    DEBUG(1,cout << "Dust_model absorbed luminosity this iteration: "  \
	  << integrate_quantity(model.absorbed(), lambda, false)*      \
	  normalization << endl;);

    cout << "Testing whether converged...  " << endl;

    // if debugging, we save the arrays for this iteration to a file.
    ofstream out;
    DEBUG(1,out.open(("ircalc_iter"+boost::lexical_cast<string>(iter)+".txt").c_str()));

    // loop over all cells to update intensity and check convergence
    // (we loop to avoid making huge temporaries)
    array_1 old_abs(nc);
    array_1 new_abs(nc);
    for (int c=0; c<nc; ++c) {
      old_abs(c) = 
	integrate_quantity((heating_intensity(c,Range::all())+
			    dust_intensity(c,Range::all()))*
			   model.total_abs_length_to_absorption(densities(c, Range::all()))*
			   4*constants::pi*volume(c)*area_factor,
			   lambda, false);
      new_abs(c) = 
	integrate_quantity(grid.intensities()(c, Range::all())*
			   model.total_abs_length_to_absorption(densities(c, Range::all()))*
			   normalization,
			   lambda, false);
    }

    // we accumulate absorption across entire grid to ensure it's not
    // bigger than the source emission.
    grid_absorption = sum(new_abs);

    DEBUG(1,								\
	  for(int c=0; c<nc; ++c) {					\
	    /* this is the cell luminosity. it should be equal to	\
	       new_abs in the PREVIOUS iteration */			\
	    const T_float cell_lum = emission.distribution().weight(c); \
	    out << c << "\t" << new_abs(c)/old_abs(c) << '\t' << new_abs(c) << '\t' << old_abs(c) << '\t' << orig_abs(c) << '\t' << cell_lum << '\t' << volume(c) << endl;});
	  
    cout << "Absorbed luminosity this iteration: " << grid_absorption << endl;
    cout << "(" << grid_absorption/current_lum << " of emitted luminosity)\n";

    if (grid_absorption > current_lum)
       cerr << "WARNING: apparent luminosity non-conservation during ray tracing by " << grid_absorption/current_lum << endl;

    if(grid_absorption > 1.10*current_lum) {
      out.close();
      cerr << "Severe energy non-conservation, exiting!" << endl;
      assert(0);
    }

    // update dust radiative intensity
    dust_intensity += grid.intensities()(tensor::i, tensor::j)*
      normalization/(4*constants::pi*volume(tensor::i)*area_factor);
     
    out.close();
    ++iter;

    // if less than tolerance of the initial luminosity is left, we
    // meet the convergence criterion
    converged = (current_lum/L_initial<tolerance);
  }
  cout << "Converged." << endl;
  ASSERT_ALL(dust_intensity<HUGE_VAL);
  ASSERT_ALL(dust_intensity==dust_intensity);

  return false;
}



#endif
	
